package com.master.chat.llm.doubao.test;


import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionContentPart;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import com.volcengine.tos.TOSV2;
import com.volcengine.tos.TOSV2ClientBuilder;
import com.volcengine.tos.TosClientException;
import com.volcengine.tos.TosServerException;
import com.volcengine.tos.comm.HttpMethod;
import com.volcengine.tos.model.object.PreSignedURLInput;
import com.volcengine.tos.model.object.PreSignedURLOutput;
import com.volcengine.tos.model.object.PutObjectInput;
import com.volcengine.tos.model.object.PutObjectOutput;

import javax.imageio.ImageIO;
import javax.imageio.ImageWriteParam;
import javax.imageio.ImageWriter;
import javax.imageio.stream.ImageOutputStream;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class VLMChatExample {

    private static final String TOS_ENDPOINT = "https://tos-cn-beijing.volces.com";
    private static final String REGION = "cn-beijing";
//    private static final String ARK_API_KEY = System.getenv("ARK_API_KEY");
//    private static final String VOLC_ACCESSKEY = System.getenv("VOLC_ACCESSKEY");
//    private static final String VOLC_SECRETKEY = System.getenv("VOLC_SECRETKEY");
    private static final String ARK_API_KEY = "60fd752f-75cb-4f50-bdbe-415a0b3fb278";
    private static final String VOLC_ACCESSKEY = "AKLTNTQ4ZjUyNWE3ZGE4NGRlOWJhZDdmOGNmYjI3NGZiYWQ";
    private static final String VOLC_SECRETKEY ="TVRreFltRmxPV1EzWXpBeU5EZzJPRGczTURZMk9UZ3lNRGd3T0dNM016QQ==";
    private static final long TOS_EXPIRE_TIME = 3600;

    public static void compressImage(String inputImagePath, String outputImagePath, long targetSizeInBytes) throws IOException {
        File inputFile = new File(inputImagePath);
        BufferedImage img = ImageIO.read(inputFile);
        float quality = targetSizeInBytes * 1.0f / inputFile.length();
        try (ImageOutputStream ios = ImageIO.createImageOutputStream(new File(outputImagePath))) {
            ImageWriter writer = ImageIO.getImageWritersByFormatName("jpeg").next();
            ImageWriteParam param = writer.getDefaultWriteParam();
            param.setCompressionMode(ImageWriteParam.MODE_EXPLICIT);
            param.setCompressionQuality(quality);
            writer.setOutput(ios);
            writer.write(null, new javax.imageio.IIOImage(img, null, null), param);
            writer.dispose();
        }
    }

    public static void uploadImageToTOS(String filePath, String bucketName, String objectKey) throws Throwable {

        TOSV2 tos = new TOSV2ClientBuilder().build(REGION, TOS_ENDPOINT, VOLC_ACCESSKEY, VOLC_SECRETKEY);

        // step 1. 上传本地图片到tos
        File file = new File(filePath);
        try (FileInputStream inputStream = new FileInputStream(file)) {
            PutObjectInput putObjectInput = new PutObjectInput().setBucket(bucketName)
                    .setKey(objectKey).setContent(inputStream).setContentLength(file.length());
            PutObjectOutput output = tos.putObject(putObjectInput);
            System.out.println("putObject succeed, object's etag is " + output.getEtag());
            System.out.println("putObject succeed, object's crc64 is " + output.getHashCrc64ecma());
        } catch (IOException e) {
            System.out.println("putObject read file failed");
            e.printStackTrace();
            throw e;
        } catch (TosClientException e) {
            // 操作失败，捕获客户端异常，一般情况是请求参数错误，此时请求并未发送
            System.out.println("putObject failed");
            System.out.println("Message: " + e.getMessage());
            if (e.getCause() != null) {
                e.getCause().printStackTrace();
            }
            throw e;
        } catch (TosServerException e) {
            // 操作失败，捕获服务端异常，可以获取到从服务端返回的详细错误信息
            System.out.println("putObject failed");
            System.out.println("StatusCode: " + e.getStatusCode());
            System.out.println("Code: " + e.getCode());
            System.out.println("Message: " + e.getMessage());
            System.out.println("RequestID: " + e.getRequestID());
            throw e;
        } catch (Throwable t) {
            // 作为兜底捕获其他异常，一般不会执行到这里
            System.out.println("putObject failed");
            System.out.println("unexpected exception, message: " + t.getMessage());
            throw t;
        }
    }

    public static String getPreSignedURL(String bucketName, String objectKey) throws Throwable {

        TOSV2 tos = new TOSV2ClientBuilder().build(REGION, TOS_ENDPOINT, VOLC_ACCESSKEY, VOLC_SECRETKEY);

        // 生成预签名链接
        try {
            PreSignedURLInput input = new PreSignedURLInput().setBucket(bucketName).setKey(objectKey)
                    .setHttpMethod(HttpMethod.GET).setExpires(TOS_EXPIRE_TIME);
            PreSignedURLOutput output = tos.preSignedURL(input);
            System.out.println("preSignedURL succeed, the signed url is " + output.getSignedUrl());
            System.out.println("preSignedURL succeed, the signed header is " + output.getSignedHeader());
            return output.getSignedUrl();
        } catch (TosClientException e) {
            // 操作失败，捕获客户端异常，一般情况是请求参数错误，此时请求并未发送
            System.out.println("preSignedURL failed");
            System.out.println("Message: " + e.getMessage());
            if (e.getCause() != null) {
                e.getCause().printStackTrace();
            }
            throw e;
        } catch (Throwable t) {
            // 作为兜底捕获其他异常，一般不会执行到这里
            System.out.println("preSignedURL failed");
            System.out.println("unexpected exception, message: " + t.getMessage());
            throw t;
        }
    }

    public static void main(String[] args) throws Throwable {
        String filePath = "D:\\system\\Desktop\\OIP-C.jpg";
        String compressedPath = "D:\\system\\Desktop\\OIP-C111.jpeg";
        String tosBucket = "chat-master";
        // 您上传图片至存储桶的路径及名称，例如上传文件至tos的 images 下并命名为 compressed_image.jpeg 文件，则填写为 images/compressed_image.jpeg
        String objectKey = "chat-master/111.jpg";
        // 1. 压缩图片
        try {
            compressImage(filePath, compressedPath, 90 * 100);
        } catch (IOException e) {
            System.out.println("compressImage failed");
            e.printStackTrace();
            throw e;
        }
        // 2. 上传TOS
        try {
            uploadImageToTOS(compressedPath, tosBucket, objectKey);
        } catch (Throwable t) {
            System.out.println("uploadImageToTOS failed");
            t.printStackTrace();
            throw t;
        }
        // 3. 生成预签名链接
        String preSignedURL = "";
        try {
            preSignedURL = getPreSignedURL(tosBucket, objectKey);
        } catch (Throwable t) {
            System.out.println("getPreSignedURL failed");
            t.printStackTrace();
            throw t;
        }

        // 4. 调用大模型Chat
        ArkService service = new ArkService(ARK_API_KEY);

        System.out.println("----- image input -----");
        final List<ChatMessage> messages = new ArrayList<>();
        final List<ChatCompletionContentPart> multiParts = new ArrayList<>();
        multiParts.add(ChatCompletionContentPart.builder().type("image_url").imageUrl(
                new ChatCompletionContentPart.ChatCompletionContentPartImageURL(preSignedURL)
        ).build());
        multiParts.add(ChatCompletionContentPart.builder().type("text").text(
                "请识别此图片"
        ).build());
        final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER)
                .multiContent(multiParts).build();
        messages.add(userMessage);

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model("doubao-1-5-thinking-vision-pro-250428")
                .messages(messages)
                .build();

        service.createChatCompletion(chatCompletionRequest).getChoices().forEach(choice -> System.out.println(choice.getMessage().getContent()));
        // shutdown service
        service.shutdownExecutor();
    }
}
